package com.aispeech.asr.apis.ssc.service;

import static org.springframework.data.mongodb.core.query.Criteria.where;

import java.util.HashMap;

import javax.annotation.PostConstruct;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.ReactiveMongoOperations;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Service;

import com.aispeech.asr.apis.ssc.constants.Progress;
import com.aispeech.asr.apis.ssc.data.domain.TaskDoc;
import com.aispeech.asr.apis.ssc.data.form.CreateTaskForm;
import com.aispeech.asr.apis.ssc.data.repository.TaskRepository;
import com.aispeech.asr.apis.ssc.data.vo.TaskDocInfo;
import com.aispeech.asr.apis.ssc.data.vo.TaskInfo;
import com.aispeech.asr.apis.ssc.metrics.Prometheus;
import com.aispeech.asr.comm.data.pojo.ProductSettingDoc;
import com.aispeech.asr.comm.data.vo.Result;
import com.aispeech.asr.comm.enums.APICode;
import com.aispeech.asr.comm.enums.ErrorCode;
import com.aispeech.asr.comm.utils.BeanUtils;

import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Mono;

@Service
@Slf4j
public class TaskService {
    @Autowired
    TaskRepository taskRepository;
    @Autowired
    ReactiveMongoOperations reactiveMongoOperations;

    private ProductSettingDoc defaultSetting;

    @PostConstruct
    private void init() {
        defaultSetting = new ProductSettingDoc();
        defaultSetting.setConcurrencies(1);
        defaultSetting.setPriority(-1);
        defaultSetting.setQueueSize(1);
    }

    public Mono<Result> create(CreateTaskForm form) {
        String appId = form.getAppId();

        try {
            TaskDoc taskDoc = TaskDoc.of(appId, BeanUtils.toMap(form));

            return taskRepository.save(taskDoc)
                    .map(saved -> {
                        log.info("task {} created by productId: {}", saved.getId(), saved.getProductId());
                        return Result.ok(
                                TaskInfo.builder()
                                        .taskId(saved.getId())
                                        .progress(saved.getProgress())
                                        .build()
                        );
                    }).doOnTerminate(() -> addPrometheus(appId, "create"));
        } catch (Exception e) {
            log.error(String.format("create task failed. productId: %s", appId), e);
            return Mono.just(Result.fail(ErrorCode.ERR_TASK_CREATE_FAILED))
                    .doOnTerminate(() -> addPrometheus(appId, "create"));
        }
    }

    private void addPrometheus(String productId, String action) {
        Prometheus.COUNTER_TASK_API_REQUESTS.labels(
                Prometheus.Env(), Prometheus.NodeName(), productId, action
        ).inc();
    }

    public Mono<Result> progress(String productId, String taskId) {
        return taskRepository.findByIdAndProductId(taskId, productId)
                .flatMap(taskDoc -> {
                    int channel = (int) taskDoc.getRequest().getOrDefault("channel", 1);
                    String audioType = (String) taskDoc.getRequest().get("audioType");

                    if (channel > 1 && "ogg_opus".equalsIgnoreCase(audioType)) {
                        return dynamicProgressBySubTasks(taskId);
                    } else {
                        return Mono.just(Result.ok(TaskInfo.builder().taskId(taskId).progress(taskDoc.getProgress()).build()));
                    }
                })
                .defaultIfEmpty(
                        Result.fail(
                                ErrorCode.ERR_TASK_LOAD_FAILED.errid,
                                String.format("Not found task: %s, productId: %s", taskId, productId)
                        )
                ).doOnTerminate(() -> addPrometheus(productId, "progress"));
    }

    private Mono<Result> dynamicProgressBySubTasks(String taskId) {
        return taskRepository.findAllByParentId(taskId)
                .collectList()
                .map(taskDocs -> {
                    int total = 0;
                    for (TaskDoc subTaskDoc : taskDocs) {
                        total += subTaskDoc.getProgress();
                    }
                    int progress = total / taskDocs.size();

                    return Result.ok(TaskInfo.builder().taskId(taskId).progress(progress).build());
                });
    }

    public Mono<Result> result(String productId, String taskId) {
        return taskRepository.findByIdAndProductId(taskId, productId)
                .flatMap(this::checkAndGenerateResult)
                .defaultIfEmpty(
                        Result.fail(
                                ErrorCode.ERR_TASK_LOAD_FAILED.errid,
                                String.format("Not found task: %s, productId: %s", taskId, productId)
                        )
                ).doOnTerminate(() -> addPrometheus(productId, "result"));
    }
    
    public Mono<Result> internalResult(String productId, String taskId) {
        return taskRepository.findById(taskId)
                .flatMap(this::checkAndGenerateInternalResult)
                .defaultIfEmpty(
                        Result.fail(
                                ErrorCode.ERR_TASK_LOAD_FAILED.errid,
                                String.format("Not found task: %s, productId: %s", taskId, productId)
                        )
                ).doOnTerminate(() -> addPrometheus(productId, "result"));
    }

    private Mono<Result> checkAndGenerateResult(TaskDoc taskDoc) {
        if (taskDoc.getProgress() < Progress.FINISHED) {
            return Mono.just(Result.fail(APICode.NOT_FINISHED));
        }

        return Mono.just(TaskDocInfo.mergeAnnotation(taskDoc));
    }
    
    private Mono<Result> checkAndGenerateInternalResult(TaskDoc taskDoc) {
        if (taskDoc.getProgress() < Progress.FINISHED) {
            return Mono.just(Result.fail(APICode.NOT_FINISHED));
        }

        return Mono.just(TaskDoc.mergeAnnotation(taskDoc));
    }
    
    public Mono<Result> doAction(String action, String taskId) {
        switch (action) {
            case "interrupt":
                return interrupt(taskId);
            case "reset":
                return reset(taskId);
            case "query":
                return query();
            default:
                return Mono.just(Result.fail(ErrorCode.ERR_PARAMS_MISSED.errid, "unsupported action: " + action));
        }
    }
    
    private Mono<Result> interrupt(String taskId) {
        Query q = Query.query(where("_id").is(taskId));
        Update u = new Update();

        u.addToSet("error", new HashMap<String, Object>(){{
            put("errid", ErrorCode.ERR_CONTROLL_INTERRUPTED.errid);
            put("error", ErrorCode.ERR_CONTROLL_INTERRUPTED.error);
        }});
        u.addToSet("progress", Progress.FINISHED);

        return reactiveMongoOperations
                .updateFirst(q, u, TaskDoc.class)
                .map(updateResult -> Result.ok());
    }
    
    private Mono<Result> reset(String taskId) {
        Query q = Query.query(where("_id").is(taskId));
        Update u = new Update();

        u.addToSet("hostname", null);
        u.addToSet("progress", Progress.IDLE);

        return reactiveMongoOperations
                .updateFirst(q, u, TaskDoc.class)
                .map(updateResult -> Result.ok());
    }
    
    private Mono<Result> query() {
        return taskRepository.findAllByProgressLessThan(Progress.FINISHED)
                .collectList()
                .map(taskDocs -> Result.ok(taskDocs));
    }

}
